#ifndef __INTERGRIDOPERATOR_H__
#define __INTERGRIDOPERATOR_H__

#include <numeric>
#include "MultiArrays.h"
#include "Vec.h"

namespace GridOp
{
    template <typename T, int N>
    class Interpolation
    {
    public:
        /**
         * @brief Interpolate between 2 arrays in a multiarray.
         *
         * @param a
         * @param l : interpolate level l-1 from level l
         * @param isCover : is cover the original value ;
         *        if false will add up to the original value
         */
        virtual void operator()(MultiArrays<T, N> &a, int l, bool isCover) const = 0;
    };

    template <typename T, int N>
    class Restriction
    {
    public:
        /**
         * @brief Restrict between 2 arrays in a multiarray.
         *
         * @param a
         * @param l : restrict level l to level l+1
         */
        virtual void operator()(MultiArrays<T, N> &a, int l) const = 0;
    };

    template <typename T, int N>
    class Linear_Interp : public Interpolation<T, N>
    {
    public:
        virtual void operator()(MultiArrays<T, N> &a, int l, bool isCover) const
        {
            if (l <= 0)
                throw std::out_of_range("Cannot perform interpolation on the finest grid!");

            int coarser_size = a.getFinestSize() >> l, finer_size = coarser_size * 2;
            auto finer_grid = a.Level(l - 1), coarser_grid = a.Level(l);

            Vec<int, N> coord;
            // the indices of neighbors in the coarser grid
            std::vector<int> coarser_neighbor_idx;

            int finer_grid_num = std::pow(finer_size + 1, N);
            for (int idx = 0; idx < finer_grid_num; ++idx)
            {
                coarser_neighbor_idx.clear();
                coarser_neighbor_idx.push_back(0);
                coord = idx2coord<N>(idx, finer_size);
                for (int i = 0; i < N; ++i)
                {
                    int num_neighbors = coarser_neighbor_idx.size();
                    if (coord[i] % 2 == 0)
                    {
                        for (int j = 0; j < num_neighbors; ++j)
                        {
                            coarser_neighbor_idx[j] += (coord[i] / 2) * std::pow(coarser_size + 1, i);
                        }
                    }
                    else if (coord[i] % 2 == 1)
                    {
                        for (int j = 0; j < num_neighbors; ++j)
                        {
                            coarser_neighbor_idx[j] += (coord[i] / 2) * std::pow(coarser_size + 1, i);
                            coarser_neighbor_idx.push_back(coarser_neighbor_idx[j] + std::pow(coarser_size + 1, i));
                        }
                    }
                }

                T new_value = 0;
                for (int coarse_idx : coarser_neighbor_idx)
                {
                    new_value += coarser_grid[coarse_idx];
                }
                new_value /= coarser_neighbor_idx.size();
                if (isCover)
                    finer_grid[idx] = std::move(new_value);
                else
                    finer_grid[idx] += std::move(new_value);
            }
        }
    };

    template <typename T, int N>
    class Quadra_Interp : public Interpolation<T, N>
    {
    public:
        virtual void operator()(MultiArrays<T, N> &a, int l, bool isCover) const
        {
            if (l <= 0)
                throw std::out_of_range("Cannot perform interpolation on the finest grid!");

            int coarser_size = a.getFinestSize() >> l, finer_size = coarser_size * 2;
            auto finer_grid = a.Level(l - 1), coarser_grid = a.Level(l);

            Vec<int, N> coord;
            std::vector<int> coarser_neighbor_idx;
            std::vector<int> coarser_neighbor_coef;

            int finer_grid_num = std::pow(finer_size + 1, N);
            for (int idx = 0; idx < finer_grid_num; ++idx)
            {
                coarser_neighbor_idx.clear();
                coarser_neighbor_coef.clear();
                coarser_neighbor_idx.push_back(0);
                coarser_neighbor_coef.push_back(1);
                coord = idx2coord<N>(idx, finer_size);
                for (int i = 0; i < N; ++i)
                {
                    int num_neighbors = coarser_neighbor_idx.size();
                    if (coord[i] % 2 == 0)
                    {
                        for (int j = 0; j < num_neighbors; ++j)
                        {
                            coarser_neighbor_idx[j] += (coord[i] / 2) * std::pow(coarser_size + 1, i);
                        }
                    }
                    else if (coord[i] % 2 == 1)
                    {
                        for (int j = 0; j < num_neighbors; ++j)
                        {
                            coarser_neighbor_idx[j] += (coord[i] / 2) * std::pow(coarser_size + 1, i);
                            coarser_neighbor_idx.push_back(coarser_neighbor_idx[j] + std::pow(coarser_size + 1, i));
                            if (coord[i] == 1)
                            {
                                coarser_neighbor_idx.push_back(coarser_neighbor_idx[j] + 2 * std::pow(coarser_size + 1, i));
                                coarser_neighbor_coef.push_back(6 * coarser_neighbor_coef[j]);
                                coarser_neighbor_coef.push_back(-coarser_neighbor_coef[j]);
                                coarser_neighbor_coef[j] *= 3;
                            }
                            else
                            {
                                coarser_neighbor_idx.push_back(coarser_neighbor_idx[j] - std::pow(coarser_size + 1, i));
                                coarser_neighbor_coef.push_back(3 * coarser_neighbor_coef[j]);
                                coarser_neighbor_coef.push_back(-coarser_neighbor_coef[j]);
                                coarser_neighbor_coef[j] *= 6;
                            }
                        }
                    }
                }

                T new_value = 0;
                for (int i = 0; i < coarser_neighbor_idx.size(); ++i)
                {
                    new_value += coarser_grid[coarser_neighbor_idx[i]] * coarser_neighbor_coef[i];
                }
                int neighbor_num = coarser_neighbor_idx.size();

                int pow_3 = 0;
                while (neighbor_num > 1)
                {
                    neighbor_num /= 3;
                    pow_3++;
                }
                new_value /= std::pow(8, pow_3);

                if (isCover)
                    finer_grid[idx] = std::move(new_value);
                else
                    finer_grid[idx] += std::move(new_value);
            }
        }
    };

    template <typename T, int N>
    class Injection : public Restriction<T, N>
    {
    public:
        virtual void operator()(MultiArrays<T, N> &a, int l) const
        {
            if (l >= a.getLevelNum())
                throw std::out_of_range("Cannot perform restriction on coasest grid!");

            int coarser_size = a.getFinestSize() >> (l + 1), finer_size = coarser_size * 2;
            auto finer_grid = a.Level(l), coarser_grid = a.Level(l + 1);
            for (int idx = 0; idx < std::pow(coarser_size + 1, N); ++idx)
            {
                // the according idx in the finer grid
                int cor_finer_idx = coord2idx<N>(idx2coord<N>(idx, coarser_size) * 2, finer_size);
                coarser_grid[idx] = finer_grid[cor_finer_idx];
            }
        }
    };

    template <typename T, int N>
    class FullWeight : public Restriction<T, N>
    {
    public:
        virtual void operator()(MultiArrays<T, N> &a, int l) const
        {
            if (l >= a.getLevelNum())
                throw std::out_of_range("Cannot perform restriction on coasest grid!");

            int coarser_size = a.getFinestSize() >> (l + 1), finer_size = coarser_size * 2;
            auto finer_grid = a.Level(l), coarser_grid = a.Level(l + 1);

            std::vector<int> finer_neighbor_idx;
            std::vector<int> finer_neighbor_coeff;
            for (int idx = 0; idx < std::pow(coarser_size + 1, N); ++idx)
            {
                int cor_finer_idx = coord2idx<N>(idx2coord<N>(idx, coarser_size) * 2, finer_size);
                finer_neighbor_idx.clear();
                finer_neighbor_coeff.clear();
                finer_neighbor_idx.push_back(cor_finer_idx);
                finer_neighbor_coeff.push_back(1);
                Vec<int, N> coord = idx2coord<N>(idx, coarser_size);
                for (int dim = 0, coord_unit = 1; dim < N; ++dim, coord_unit *= (finer_size + 1))
                {
                    if (coord[dim] == 0 || coord[dim] == coarser_size)
                    {
                        coarser_grid[idx] = finer_grid[cor_finer_idx];
                        continue;
                    }
                    int num_neighbors = finer_neighbor_idx.size();
                    for (int i = 0; i < num_neighbors; ++i)
                    {
                        finer_neighbor_idx.push_back(finer_neighbor_idx[i] + coord_unit);
                        finer_neighbor_idx.push_back(finer_neighbor_idx[i] - coord_unit);
                        finer_neighbor_coeff.push_back(finer_neighbor_coeff[i]);
                        finer_neighbor_coeff.push_back(finer_neighbor_coeff[i]);
                        finer_neighbor_coeff[i] *= 2;
                    }
                }

                coarser_grid[idx] = 0;
                int num_neighbors = finer_neighbor_idx.size();
                for (int i = 0; i < num_neighbors; ++i)
                {
                    coarser_grid[idx] += finer_grid[finer_neighbor_idx[i]] * finer_neighbor_coeff[i];
                }

                for (; num_neighbors > 1; num_neighbors /= 3)
                {
                    coarser_grid[idx] /= 4;
                }
            }
        }
    };
}

#endif // __INTERGRIDOPERATOR_H__